import cv2
import face_recognition
from frame_rates import FrameRates
from play import ffplay
import numpy as np
from known_faces import known_faces, known_faces_names

mode = 'export'
if mode == 'export':
    from process_result import ffmpeg_res
    from process_capture import ffmpeg_cap

fr = FrameRates()
# 人脸误差阈值
threshold = 0.6
cap = cv2.VideoCapture(0)
while True:
    ret, frame = cap.read()
    if mode == 'export':
        # 保存原始视频
        ffmpeg_cap.stdin.write(frame.tostring())
        ffmpeg_cap.stdin.flush()

    fr.record()
    gray_face = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # 识别图片中的脸部（可能存在多个脸）
    face_locations = face_recognition.face_locations(gray_face)
    # 遍历人脸位置信息
    # face_locations = []
    for top, right, bottom, left in face_locations:
        # 对人脸画框
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)
        # 识别这张人脸对应的姓名
        face_img = frame[top:bottom, left:right]
        # 提取人脸编码
        current_face_encoding = face_recognition.face_encodings(face_img)
        if current_face_encoding:
            # 和现有人脸比较相似度
            similar_matrix = face_recognition.compare_faces(known_faces, current_face_encoding[0], tolerance=threshold)
            name_idx = np.array(similar_matrix).searchsorted(True)
            if name_idx < len(known_faces_names):
                # 提取面部对应的姓名
                face_name = known_faces_names[name_idx]
                # 将姓名添加到图片上
                cv2.putText(frame, face_name, (left, top - 5), cv2.FONT_HERSHEY_DUPLEX, 0.8, (0, 255, 0), 1)

    # 计算帧率
    cv2.putText(frame, f"fps: {fr.get_fps()}", (30, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (240, 240, 240), 1)
    # 打印帧率
    print(fr.get_fps())
    # ffplay 显示视频
    ffplay.stdin.write(frame.tostring())
    ffplay.stdin.flush()

    if mode == 'export':
        # 保存结果视频
        ffmpeg_res.stdin.write(frame.tostring())
        ffmpeg_res.stdin.flush()
    if cv2.waitKey(100) & 0xff == ord('q'):
        break
cap.release()
